/*---------------------------------------------------------------------------*\
|* This file is part of the koala project with MIT License.                  *|
|* Copyright (c) James <zhuguangxiang@gmail.com>                             *|
\*--------------------------------------------------------------------------**/

#ifndef _KOALA_IR_H_
#define _KOALA_IR_H_

#include "util/hashmap.h"
#include "util/list.h"
#include "util/vector.h"

#ifdef __cplusplus
extern "C" {
#endif

/* Koala Language intermediate Representation
 * It is a high level representation of Koala language with some medium level
 * opcodes. https://cs.lmu.edu/~ray/notes/ir/
 */

/* Koala language supports i8, i16, i32, i64, f32, f64 and so on. But all i8 and
 * i16 values are casted into i32, then do computing. It's like Java language.
 * It sounds reasonable, because i8 and i16 are too short and often overflow.
 * Koala also does not support unsigned integer, all integers are signed.
 */
typedef enum _KlrTypeKind {
    KLR_TYPE_VOID = 1,
    KLR_TYPE_INT8,
    KLR_TYPE_INT16,
    KLR_TYPE_INT32,
    KLR_TYPE_INT64,
    KLR_TYPE_FLOAT32,
    KLR_TYPE_FLOAT64,
    KLR_TYPE_BOOL,
    KLR_TYPE_CHAR,
    KLR_TYPE_STRING,
    KLR_TYPE_ARRAY,
    KLR_TYPE_MAP,
    KLR_TYPE_PROTO,
    KLR_TYPE_FUNC,
    KLR_TYPE_CLASS,
    KLR_TYPE_INTF,
    KLR_TYPE_ENUM,
    KLR_TYPE_CLOSURE,
} KlrTypeKind;

typedef struct _KlrType {
    KlrTypeKind kind;
} KlrType;

typedef struct _KlrProtoType {
    KlrTypeKind kind;
    Vector *params;
    KlrType *ret;
} KlrProtoType;

typedef struct _KlrKlassType {
    KlrTypeKind kind;
    char *name;
    Vector fields;
    Vector methods;
} KlrKlassType;

typedef struct _KlrIntfType {
    KlrTypeKind kind;
    char *name;
    Vector intfs;
} KlrIntfType;

typedef struct _KlrEnumType {
    KlrTypeKind kind;
    char *name;
} KlrEnumType;

KlrType *klr_type_void(void);
KlrType *klr_type_int8(void);
KlrType *klr_type_int16(void);
KlrType *klr_type_int32(void);
KlrType *klr_type_int64(void);
KlrType *klr_type_float32(void);
KlrType *klr_type_float64(void);
KlrType *klr_type_bool(void);
KlrType *klr_type_char(void);
KlrType *klr_type_string(void);
KlrType *klr_type_proto(KlrType *ret, KlrType *params[], int num_params);

typedef enum _KlrValueKind {
    KLR_VALUE_CONST = 1,
    KLR_VALUE_GLOBAL,
    KLR_VALUE_FUNC,
    KLR_VALUE_BLOCK,
    KLR_VALUE_PARAM,
    KLR_VALUE_LOCAL,
    KLR_VALUE_INSN,
} KlrValueKind;

/* clang-format off */
#define KLR_VALUE_HEAD KlrValueKind kind; KlrType *type; \
    List use_list; char *name; int tag; int vreg;
/* clang-format on */

typedef struct _KlrValue {
    KLR_VALUE_HEAD
} KlrValue;

typedef struct _KlrConst {
    KLR_VALUE_HEAD
    union {
        int8_t i8val;
        int16_t i16val;
        int32_t i32val;
        int64_t i64val;
        float f32val;
        double f64val;
        int bval;
        uint32_t cval;
        char *sval;
    };
} KlrConst;

KlrValue *klr_const_int8(int8_t v);
KlrValue *klr_const_int16(int16_t v);
KlrValue *klr_const_int32(int32_t v);
KlrValue *klr_const_int64(int64_t v);
KlrValue *klr_const_float32(float v);
KlrValue *klr_const_float64(double v);
KlrValue *klr_const_bool(int v);
KlrValue *klr_const_char(uint32_t v);
KlrValue *klr_const_string(char *v);

typedef struct _KlrGlobal {
    KLR_VALUE_HEAD
} KlrGlobal;

/* local register variable */
typedef struct _KlrLocal {
    KLR_VALUE_HEAD
    /* link in bb */
    List bb_link;
    /* ->bb */
    struct _KlrBasicBlock *bb;
} KlrLocal;

typedef struct _KlrParam {
    KLR_VALUE_HEAD
} KlrParam;

typedef struct _KlrFunc {
    KLR_VALUE_HEAD
    /* value tag */
    int val_tag;
    /* basic block tag */
    int bb_tag;
    /* number of instructions */
    int num_insns;
    /* number of registers */
    int vregs;
    /* instructions list */
    List insn_list;
    /* basic block list */
    List bb_list;
    /* all edges */
    List edge_list;
    /* params(value, use only) */
    Vector params;
    /* locals(reg, both def/use) */
    Vector locals;
    /* start basic block */
    struct _KlrBasicBlock *sbb;
    /* end basic block */
    struct _KlrBasicBlock *ebb;
} KlrFunc;

typedef struct _KlrBasicBlock {
    KLR_VALUE_HEAD
    /* linked in KlrFunc */
    List link;
    /* ->KlrFunc(parent) */
    KlrFunc *func;

    /* block has branch flag */
    int has_branch;
    /* block has return flag */
    int has_ret;
    /* visited? */
    int visited;

    /* number of instructions */
    int num_insns;
    /* number of in-edges */
    int num_inedges;
    /* number of out-edges */
    int num_outedges;

    /* locals defined in this block */
    List local_list;
    /* instructions */
    List insn_list;
    /* in edges(predecessors) */
    List in_edges;
    /* out edges(successors) */
    List out_edges;
} KlrBasicBlock;

typedef struct _KlrEdge {
    /* linked in KlrFunc */
    List link;
    /* src block */
    KlrBasicBlock *src;
    /* dest block */
    KlrBasicBlock *dst;
    /* add to dest block */
    List in_link;
    /* add to src block */
    List out_link;
} KlrEdge;

typedef struct _KlrUse {
    /* parent(use_list) */
    KlrValue *parent;
    /* link in use_list */
    List use_link;
    /* self instruction */
    struct _KlrInsn *insn;
} KlrUse;

#define use_foreach(use, val, closure) \
    list_foreach(use, use_link, &(val)->use_list, closure)

typedef enum _KlrOperKind {
    KLR_OPER_CONST = 1,
    KLR_OPER_GLOBAL,
    KLR_OPER_FUNC,
    KLR_OPER_BLOCK,
    KLR_OPER_LOCAL,
    KLR_OPER_REG,
} KlrOperKind;

typedef struct _KlrOper {
    KlrOperKind kind;
    union {
        KlrConst *konst;
        KlrUse *use;
    };
} KlrOper;

/* instruction opcodes */
typedef enum _KlrInsnKind {
    KLR_INSN_SET = 1,
    KLR_INSN_GET,

    KLR_INSN_ADD,
    KLR_INSN_SUB,
    KLR_INSN_MUL,
    KLR_INSN_DIV,
    KLR_INSN_MOD,

    KLR_INSN_AND,
    KLR_INSN_OR,
    KLR_INSN_XOR,
    KLR_INSN_SHL,
    KLR_INSN_SHR,
    KLR_INSN_USHR,

    KLR_INSN_CMPEQ,
    KLR_INSN_CMPNE,
    KLR_INSN_CMPGT,
    KLR_INSN_CMPGE,
    KLR_INSN_CMPLT,
    KLR_INSN_CMPLE,

    KLR_INSN_LAND,
    KLR_INSN_LOR,
    KLR_INSN_LNOT,

    KLR_INSN_NEW,
    KLR_INSN_NEW_ARRAY,
    KLR_INSN_NEW_MAP,

    KLR_INSN_ARRAY_GET,
    KLR_INSN_ARRAY_SET,
    KLR_INSN_ARRAY_LEN,

    KLR_INSN_MAP_GET,
    KLR_INSN_MAP_SET,

    KLR_INSN_CALL,
    KLR_INSN_ICALL,
    KLR_INSN_VCALL,

    KLR_INSN_GET_FIELD,
    KLR_INSN_SET_FIELD,

    KLR_INSN_GET_SUBSCR,
    KLR_INSN_SET_SUBSCR,

    KLR_INSN_JMP,
    KLR_INSN_COND_JMP,
    KLR_INSN_RET,
    KLR_INSN_RET_VOID,

} KlrInsnKind;

typedef struct _KlrInsn {
    KLR_VALUE_HEAD
    /* linear position */
    int pos;
    /* opcode */
    KlrInsnKind opcode;
    /* number of operands */
    int num_opers;
    /* link in bb */
    List bb_link;
    /* ->bb */
    KlrBasicBlock *bb;
    /* link in KlrFunc */
    List fn_link;
    /* ->KlrFunc */
    KlrFunc *fn;
    /* operands */
    KlrOper opers[0];
} KlrInsn;

typedef struct _KlrModule {
    /* module name */
    char *name;
    /* all symbols */
    HashMap stbl;
    /* variables */
    Vector variables;
    /* functions */
    Vector functions;
    /* types */
    Vector types;
    /* __init__ function */
    KlrFunc *init;
} KlrModule;

typedef struct _KlrBuilder {
    KlrBasicBlock *bb;
    List *it;
} KlrBuilder;

/* basic blocks & edges */

/* append a basic block to the end of a function */
KlrBasicBlock *klr_append_block(KlrValue *fn_val, char *label);

/* add a basic block after 'bb' */
KlrBasicBlock *klr_add_block(KlrBasicBlock *bb, char *label);

/* add a basic block before 'bb' */
KlrBasicBlock *klr_add_block_before(KlrBasicBlock *bb, char *label);

/* delete a basic block */
void klr_delete_block(KlrBasicBlock *bb);

/* add an edge */
void klr_link_age(KlrBasicBlock *src, KlrBasicBlock *dst);

/* edge-out iteration */
#define edge_out_foreach(edge, bb, closure) \
    list_foreach(edge, out_link, &(bb)->out_edges, closure)

/* edge-in iteration */
#define edge_in_foreach(edge, bb, closure) \
    list_foreach(edge, in_link, &(bb)->in_edges, closure)

#define edge_out_empty(bb) list_empty(&(bb)->out_edges)
#define edge_out_first(bb) list_first(&(bb)->out_edges, KlrEdge, out_link)
#define edge_out_last(bb)  list_last(&(bb)->out_edges, KlrEdge, out_link)

#define edge_in_empty(bb) list_empty(&(bb)->in_edges)
#define edge_in_first(bb) list_first(&(bb)->in_edges, KlrEdge, in_link)
#define edge_in_last(bb)  list_last(&(bb)->in_edges, KlrEdge, in_link)

/* basic block iteration */
#define basic_block_foreach(bb, fn, closure) \
    list_foreach(bb, link, &(fn)->bb_list, closure)

/* set builder at head */
static inline void klr_builder_head(KlrBuilder *bldr, KlrBasicBlock *bb)
{
    bldr->bb = bb;
    bldr->it = &bb->insn_list;
}

/* set builder at and */
static inline void klr_builder_end(KlrBuilder *bldr, KlrBasicBlock *bb)
{
    bldr->bb = bb;
    bldr->it = bb->insn_list.prev;
}

/* set builder at 'insn' */
static inline void klr_builder_at(KlrBuilder *bldr, KlrInsn *insn)
{
    bldr->it = &insn->bb_link;
}

/* set builder before 'insn' */
static inline void klr_builder_before(KlrBuilder *bldr, KlrInsn *insn)
{
    bldr->it = insn->bb_link.prev;
}

/* instructions */

KlrValue *klr_get_param(KlrValue *fn_val, int index);
KlrValue *klr_add_local(KlrBuilder *bldr, KlrType *ty, char *name);
void klr_build_set(KlrBuilder *bldr, KlrValue *var, KlrValue *val);
KlrValue *klr_build_get(KlrBuilder *bldr, KlrValue *var);

KlrValue *klr_build_binary(KlrBuilder *bldr, KlrValue *lhs, KlrValue *rhs,
                           KlrInsnKind op, char *name);

#define klr_build_add(bldr, lhs, rhs, name) \
    klr_build_binary(bldr, lhs, rhs, KLR_INSN_ADD, name)

#define klr_build_sub(bldr, lhs, rhs, name) \
    klr_build_binary(bldr, lhs, rhs, KLR_INSN_SUB, name)

#define klr_build_mul(bldr, lhs, rhs, name) \
    klr_build_binary(bldr, lhs, rhs, KLR_INSN_MUL, name)

#define klr_build_div(bldr, lhs, rhs, name) \
    klr_build_binary(bldr, lhs, rhs, KLR_INSN_DIV, name)

#define klr_build_mod(bldr, lhs, rhs, name) \
    klr_build_binary(bldr, lhs, rhs, KLR_INSN_MOD, name)

#define klr_build_cmpeq(bldr, lhs, rhs, name) \
    klr_build_binary(bldr, lhs, rhs, KLR_INSN_CMPEQ, name)

#define klr_build_cmpne(bldr, lhs, rhs, name) \
    klr_build_binary(bldr, lhs, rhs, KLR_INSN_CMPNE, name)

#define klr_build_cmpgt(bldr, lhs, rhs, name) \
    klr_build_binary(bldr, lhs, rhs, KLR_INSN_CMPGT, name)

#define klr_build_cmpge(bldr, lhs, rhs, name) \
    klr_build_binary(bldr, lhs, rhs, KLR_INSN_CMPGE, name)

#define klr_build_cmplt(bldr, lhs, rhs, name) \
    klr_build_binary(bldr, lhs, rhs, KLR_INSN_CMPLT, name)

#define klr_build_cmple(bldr, lhs, rhs, name) \
    klr_build_binary(bldr, lhs, rhs, KLR_INSN_CMPLE, name)

void klr_build_jmp(KlrBuilder *bldr, KlrBasicBlock *dst);
void klr_build_condjmp(KlrBuilder *bldr, KlrValue *cond, KlrBasicBlock *_then,
                       KlrBasicBlock *_else);
void klr_build_ret(KlrBuilder *bldr, KlrValue *ret);
void klr_build_ret_void(KlrBuilder *bldr);

/* instruction iteration */
#define insn_foreach(insn, bb, closure) \
    list_foreach(insn, bb_link, &(bb)->insn_list, closure)

/* modules & functions */

static inline void klr_set_name(KlrValue *val, char *name)
{
    val->name = name ? name : "";
}

KlrModule *klr_create_module(char *name);
void klr_destroy_module(KlrModule *m);

KlrValue *klr_add_func(KlrModule *m, KlrType *ty, char *name);
KlrValue *klr_add_global(KlrModule *m, KlrType *ty, char *name);

#ifdef __cplusplus
}
#endif

#endif /* _KOALA_IR_H_ */
